#pragma once

#include "../commons/token.hpp"
#include "../commons/utils.hpp"
#include <algorithm>
#include <deque>
#include <iostream>
#include <memory>
#include <set>
#include <string>
#include <unordered_map>
#include <vector>
#include <stack>
#include <map>

using std::cerr;
using std::cout;
using std::deque;
using std::endl;
using std::fstream;
using std::make_shared;
using std::map;
using std::set;
using std::shared_ptr;
using std::stack;
using std::string;
using std::stringstream;
using std::unordered_map;
using std::vector;

/**
 * The LR(1) parser
 */ 
class MyParser {
public:
    enum class ActionType { //分析表动作类型
        UNDEFINED,
        ACCEPT,
        SHIFT,
        REDUCE,
        CONFLICT,
        ERROR
    };                   
    class RuleType; // 规则类型
    class LRItem;   // 一个LR1项目
    class Action;   // 分析表动作
    class ParseTreeNode; // 分析树节点
    using PtrAction = shared_ptr<Action>;
    using PtrParseTreeNode = shared_ptr<ParseTreeNode>;

public:
    /**
     * Build parser from syntax defined in CSV file, it's format is
     * rule1, token1 token2 token3
     * rule2, token1 token2 token3
     * @param iss the input stream
     */ 
    void buildParserFromSyntaxDef(std::istream &iss);
    /**
     * Build parser from parse table
     * @param csvin the input stream of pasre table
     * @param syntaxin the input stream of syntax definition file
     */
    void buildParserFromCsvParseTable(std::istream &csvin, std::istream &syntaxin);
    /**
     * Parse from token sequence
     * @param tokens token sequence
     * @return parse tree root
     */ 
    PtrParseTreeNode parse(const vector<Token> &tokens);
    void printParseTable(std::ostream &oss) const;
    void printParseTree(std::ostream &oss, const PtrParseTreeNode &root,
                        string indent) const;
    void printSyntax(std::ostream &oss) const;

public:
    /**
     * 规则类型
     */
    class RuleType {
    public:
        string name;     // lhs在symbol里的名字
        int index;       // 该规则在rules数组中的下标
        int lhs;         // 该规则名称在symbols数组中的下标
        vector<int> rhs; // 该规则包含的子规则和终结符在symbols数组中的下标
        RuleType();
        RuleType(int index, string name);
        string getName() const;
        bool operator==(const RuleType &t) const;
        bool operator!=(const RuleType &t) const;
        bool operator<(const RuleType &item) const;
    };

    /**
     * LR1项目
     */
    class LRItem {
    public:
        RuleType rule;
        int position;
        int lookahead;
        LRItem() = default;
        LRItem(RuleType rule, int position, int lookahead);
        int getExpectedSymbol();
        bool operator==(const LRItem &item) const;
        bool operator!=(const LRItem &item) const;
        bool operator<(const LRItem &item) const;
    };

    /**
     * 分析表动作
     */
    class Action {
    public:
        ActionType atype;
        int state;
        RuleType rule;
        vector<PtrAction> actions;
        Action();
        Action(bool);
        Action(int state);
        Action(RuleType rule);
        ~Action() {}
        void addAction(const Action &a);
        friend std::ostream &operator<<(std::ostream &oss, const Action &a);
        bool operator==(const Action &a) const;
        bool operator!=(const Action &b) const;
        void copyFields(const Action &a);
    };

    /**
     * 分析树节点
     */
    class ParseTreeNode {
    public:
        union {
            RuleType ruleT;
            Token token;
        };
        deque<PtrParseTreeNode> children;

        bool isTerminal();
        ParseTreeNode();
        ParseTreeNode(RuleType ruleT, deque<PtrParseTreeNode> children);
        ParseTreeNode(Token token);
        ~ParseTreeNode();
    };

private:
    const int _startSymbol = 0;     // 分析栈内起始符号
    vector<string> _symbols;        // 语法符号
    int _symbolsTerminalOffset;     // _symbols[0, offset) 为非终结符，后面的为终结符
    int _endMarker;                 // _symbols数组长度
    const string _ENDCHAR = "$";           // 终结符 $
    unordered_map<int, set<int>> first;    // first集
    vector<RuleType> _rules;               // 规则
    vector<set<LRItem>> _collection;    // 状态机状态，_collection[i]表示 i 状态下的项目集合
    vector<vector<Action>> _parseTable;    // 动作
    unordered_map<string, int> _tokenName2Index;            // 将token名转换成index
    bool _parseTableHasConflict;                            // LR表有冲突

private:
    void clear();
    void readSyntax(std::istream &iss);
    bool isNonTerminal(int);
    bool isTerminal(int);
    bool isTerminalOrEndMarker(int);
    bool setsEqual(set<LRItem>,set<LRItem>);
    void getFirst(int);
    set<LRItem> closure(set<LRItem>);
    set<LRItem> goTo(set<LRItem>, int);
    void printState(int);
    void createCollection();
    void addActionToParseTable(int, int, Action);
    void createParseTable();
    int token2Symbol(const Token &);

public:
    vector<string> getSymbols() const;
    vector<vector<Action>> getParseTable() const;
    bool hasConflict() const { return _parseTableHasConflict; }
};

void MyParser::clear() {
    _symbols.clear();
    _rules.clear();
    _collection.clear();
    _parseTable.clear();
    _tokenName2Index.clear();
    _parseTableHasConflict = false;
}

inline bool MyParser::isNonTerminal(int symbol) {
    return symbol < _symbolsTerminalOffset && symbol > 0;
}

inline bool MyParser::isTerminal(int symbol) {
    return symbol >= _symbolsTerminalOffset && symbol < _symbols.size() - 1;
}

inline bool MyParser::isTerminalOrEndMarker(int symbol) {
    return symbol >= _symbolsTerminalOffset && symbol < _symbols.size();
}

bool MyParser::setsEqual(set<LRItem> s1, set<LRItem> s2) {
    return s1 == s2;
    // if (s1.size() != s2.size()) {
    //     return false;
    // }
    // sort(s1.begin(), s1.end());
    // sort(s2.begin(), s2.end());
    // for (int i = 0; i < s1.size(); ++i) {
    //     if (s1[i] != s2[i])
    //         return false;
    // }
    // return true;
}

void MyParser::getFirst(int symbol) {
    if (first.count(symbol)) {
        return;
    }
    if (isTerminalOrEndMarker(symbol)) {
        first.insert({symbol, {symbol}});
        return;
    }
    set<int> f;
    for (auto i : _rules) {
        if (i.lhs == symbol && !i.rhs.empty() && i.rhs[0] != symbol) {
            getFirst(i.rhs[0]);
            f.insert((first[i.rhs[0]]).begin(), first[i.rhs[0]].end());
        }
    }
    first.insert({symbol, f});
};

set<MyParser::LRItem> MyParser::closure(set<MyParser::LRItem> s) {
    // cout << "Closure: ";
    // for(auto i: s){
    //     cout << _symbols[i.rule.lhs) << ", ";
    // }
    // cout << endl;
    set<MyParser::LRItem> closureSet = s;
    bool newItemAdded = true;
    while (newItemAdded) {
        newItemAdded = false;
        for (auto item: closureSet) {
            const int symbol = item.getExpectedSymbol();
            for (auto rule : _rules) {
                if (rule.lhs == symbol) { // 找到拥有该产生式的规则
                    set<int> fst =
                        first[item.position + 1 < item.rule.rhs.size()
                                  ? item.rule.rhs[item.position + 1]
                                  : item.lookahead];
                    for (auto k : fst) {
                        LRItem newItem(rule, 0, k);
                        if (!contains(closureSet, newItem)) {
                            closureSet.insert(newItem);
                            newItemAdded = true;
                        }
                    }
                }
            }
        }
    }
    return closureSet; // TODO many set
};

set<MyParser::LRItem> MyParser::goTo(set<MyParser::LRItem> s, int symbol) {
    set<MyParser::LRItem> gotoSet;
    for (auto i : s) {
        if (i.position < i.rule.rhs.size() && i.rule.rhs[i.position] == symbol) {
            gotoSet.insert(LRItem(i.rule, i.position + 1, i.lookahead));
            // cout << _symbols[gotoSet.back().rule.lhs) << endl;
        }
    }
    // cout << "----\n";
    return closure(gotoSet);
}

void MyParser::printState(int i) {
    cout << "State " << i << "\n";
    for (auto j : _collection[i]) {
        cout << _symbols[j.rule.lhs] << " -> ";
        for (int k = 0; k < j.rule.rhs.size(); ++k) {
            if (k == j.position) {
                cout << ".";
            }
            cout << _symbols[j.rule.rhs[k]] << " ";
        }
        if (j.position == j.rule.rhs.size())
            cout << ".";
        cout << " @ " << _symbols[j.lookahead];
        cout << endl;
    }
    cout << endl << endl;
}

void MyParser::createCollection() {
    cout << "[Info] Syntax rules read successfully, start building LR(1) collections...\n";
    LRItem initialItem = LRItem(_rules[0], 0, _endMarker);
    _collection.push_back(closure({initialItem}));
    bool newSetAdded = true;
    while (newSetAdded) {
        newSetAdded = false;
        cout << _collection.size() << ", ";
        for (int i = 0; i < _collection.size(); ++i) {
            for (int symbol = 1; symbol < _symbols.size() - 1; ++symbol) {
                set<LRItem> nextState = goTo(_collection[i], symbol);
                if (nextState.size() != 0) {
                    bool newSet = true;
                    for (int j = 0; j < _collection.size(); ++j) {
                        if (setsEqual(_collection[j], nextState)) {
                            newSet = false;
                            break;
                        }
                    }
                    if (newSet) {
                        _collection.push_back(nextState);
                        newSetAdded = true;
                        // for (int j = 0; j < _collection.size(); ++j) {
                        //     _fromToLookup.insert({{j, symbol}, (int)_collection.size() - 1});
                        // }
                        // cout <<"From "<< i<<  " symbol: " << _symbols[symbol) << endl;
                        // printState(_collection.size() - 1);
                    }
                }
            }
        }
    }
}

void MyParser::addActionToParseTable(int state, int symbol, Action action) {
    if (_parseTable[state][symbol].atype == ActionType::UNDEFINED) { // 无冲突
        _parseTable[state][symbol] = action;
    } else if (_parseTable[state][symbol].atype == ActionType::CONFLICT) { // 有冲突
        bool newAction = true;
        for (int i = 0; i < _parseTable[state][symbol].actions.size(); ++i) {
            if (*_parseTable[state][symbol].actions[i] == action) {
                newAction = false;
                break;
            }
        }
        if (newAction) {
            _parseTable[state][symbol].addAction(action);
        }
    } else if (_parseTable[state][symbol] != action) { // 有冲突
        _parseTable[state][symbol].addAction(action);
        _parseTableHasConflict = true;
    }
}

void MyParser::createParseTable() {
    _parseTable =
        vector<vector<Action>>(_collection.size(), vector<Action>(_symbols.size()));
    _parseTableHasConflict = false;
    for (int i = 0; i < _collection.size(); ++i) {
        cout << i << ", ";
        for (auto item: _collection[i]) { // 遍历每一个项目查看它将要接受的符号
            const int symbol = item.getExpectedSymbol();
            if (symbol == -1) { // 到尾了，应reduce或accept
                if (item.rule.lhs != _startSymbol) { // reduce
                    addActionToParseTable(i, item.lookahead,
                                          Action(item.rule));
                } else if (item.lookahead == _endMarker) { // accept
                    _parseTable[i][_endMarker] = Action(true);
                }
            }
            if (isTerminal(symbol)) { // shift
                set<LRItem> gotoSet = goTo(_collection[i], symbol);
                for (int k = 0; k < _collection.size(); ++k) {
                    if (setsEqual(_collection[k], gotoSet)) {
                        addActionToParseTable(i, symbol, Action(k));
                    }
                }
            }
        }
        for (int symbol = 1; symbol < _symbolsTerminalOffset; ++symbol) { // 打GOTO表
            set<LRItem> gotoSet = goTo(_collection[i], symbol);
            for (int k = 0; k < _collection.size(); ++k) {
                if (setsEqual(_collection[k], gotoSet)) {
                    _parseTable[i][symbol] = k;
                }
            }
        }
    }
}

void MyParser::readSyntax(std::istream &iss){
    vector<vector<string>> rulesSymbols;
    vector<string> lines;
    string tmpStr;
    while (std::getline(iss, tmpStr)) {
        mytrim_copy(tmpStr);
        if (tmpStr.length() != 0) {
            tmpStr = myreplace_copy(tmpStr, ",", " ");
            lines.push_back(mytrim_copy(tmpStr));
        }
    }
    for (int i = 0; i < lines.size(); ++i) {
        rulesSymbols.push_back(mysplit(lines[i], " "));
        if (rulesSymbols.back().empty())
            rulesSymbols.pop_back();
        if (!contains(_symbols, rulesSymbols[i][0])) {
            _symbols.push_back(rulesSymbols[i][0]);
        }
    }
    _symbolsTerminalOffset = _symbols.size();
    for (auto i : rulesSymbols) {
        for (auto j : i) {
            if (!contains(_symbols, j)) {
                _symbols.push_back(j);
            }
        }
    }
    _endMarker = _symbols.size();
    _symbols.push_back(_ENDCHAR);
    for (int i = 0; i < rulesSymbols.size(); ++i) {
        _rules.push_back(RuleType(i, rulesSymbols[i][0]));
        _rules[i].lhs = indexOf(_symbols, rulesSymbols[i][0]);
        for (int j = 1; j < rulesSymbols[i].size(); ++j) {
            _rules[i].rhs.push_back(indexOf(_symbols, rulesSymbols[i][j]));
        }
    }
}

void MyParser::buildParserFromSyntaxDef(std::istream &ss) {
    cout << "[Info] Start building parser, wait patiently...\n";
    clock_t st = clock();
    clear();
    readSyntax(ss);
    for (int i = 0; i < _symbols.size(); ++i) {
        getFirst(i);
    }
#ifdef __DEBUG
    for (auto i : first) {
        cout << "First of " << _symbols[i.first] << ": ";
        for (auto j : i.second) {
            cout << _symbols[j] << " ";
        }
        cout << endl;
    }
#endif
    createCollection();
    cout << "[Info] Collections built successfully, creating parse table...\n";
    createParseTable();
    clock_t ed = clock();
    cout << "[Info] Parser created in " << ed - st << " ms\n";
    cout << "[Info] Parse table size: " << _parseTable.size() << "*"
         << _parseTable[0].size() << endl;
}

void MyParser::buildParserFromCsvParseTable(std::istream &csvin, std::istream &syntaxin) {
    clock_t st = clock();
    clear();
    readSyntax(syntaxin);
    string line;
    std::getline(csvin, line); // 表头
    while (std::getline(csvin, line)) {
        auto actions = mysplit(line, ",");
        _parseTable.push_back(vector<Action>(_symbols.size()));
        for (int i = 1; i < actions.size(); ++i) {
            if (myStartsWith(actions[i], "s")) {
                int toState = std::stoi(actions[i].substr(1));
                _parseTable.back()[i] = Action(toState);
            } else if (myStartsWith(actions[i], "r")) {
                int ruleIdx = std::stoi(actions[i].substr(1));
                _parseTable.back()[i] = Action(_rules[ruleIdx]);
            } else if (myStartsWith(actions[i], "ACC")) {
                _parseTable.back()[i] = Action(true);
            }
        }
    }
    clock_t ed = clock();
    cout << "[Info] Parser built in " << ed - st << " ms.\n";
}

int MyParser::token2Symbol(const Token &token) {
    string tokenName = token.tokenName;
    if (_tokenName2Index.count(tokenName)) {
        return _tokenName2Index[tokenName];
    } else {
        int index = indexOf(_symbols, tokenName);
        if (index == -1) {
            cout << "[Error] Token \'" << tokenName
                 << "\' did not exist in grammar definition, "
                 << "please check your grammar definition file.\n"
                 << "Parser terminated.\n";
            throw ParseException(token.endPos.line, token.endPos.col, token.word);
        }
        _tokenName2Index[tokenName] = index;
        return index;
    }
}

vector<string> MyParser::getSymbols() const { return _symbols; }
vector<vector<MyParser::Action>> MyParser::getParseTable() const { return _parseTable; }

MyParser::PtrParseTreeNode MyParser::parse(const vector<Token> &input) {
    if (_parseTableHasConflict) {
        // _parseSteps.push_back(ParseStepLog({},{},{},
        //     Action(ParseException(0,0,"Conflict in parse table"))));
        cerr << "Conflict in parse table!\nParser terminated.\n";
        throw ParseException("Conflict in parse table.");
    }
    stack<int> stateStack;
    stack<PtrParseTreeNode> symbolStack;
    Action action;
    stateStack.push(0);
    vector<Token>::const_iterator inputIter = input.begin();
    while (true) {
        if (inputIter == input.end()) {
            action = _parseTable[stateStack.top()][_endMarker];
        } else {
            action = _parseTable[stateStack.top()][token2Symbol(*inputIter)];
        }
        if (action.atype == ActionType::UNDEFINED) {
            cout << "[Error] line " << inputIter->startPos.line << "("
                 << inputIter->startPos.col << "): " << inputIter->tokenName
                 << ": \'" << inputIter->word <<"\'\nSyntax error!";
            throw ParseException(inputIter->startPos.line, inputIter->startPos.col,
                                 inputIter->word);
        }
        if (action.atype == ActionType::SHIFT) {
            symbolStack.push(make_shared<ParseTreeNode>(*inputIter));
            inputIter++;
            stateStack.push(action.state);
        } else if (action.atype == ActionType::REDUCE) {
            deque<PtrParseTreeNode> children; // 如果规约，则把所有归约的字符加进这里
            for (int i = 0; i < action.rule.rhs.size(); ++i) {
                children.push_front(symbolStack.top());
                symbolStack.pop();
                stateStack.pop();
            }
            PtrParseTreeNode node = make_shared<ParseTreeNode>(action.rule, children);
            symbolStack.push(node);
            // GOTO
            stateStack.push(_parseTable[stateStack.top()][action.rule.lhs].state);
        } else if (action.atype == ActionType::ACCEPT) {
            break;
        } else {
            cout << "[Error] line " << inputIter->startPos.line << "("
                 << inputIter->startPos.col << "): " << inputIter->tokenName << ": \'"
                 << inputIter->word << "\'\nSyntax error!";
            throw ParseException(inputIter->startPos.line, inputIter->startPos.col,
                                 inputIter->word);
        }
    }
    return symbolStack.top();
}



void MyParser::printParseTree(std::ostream &oss, const MyParser::PtrParseTreeNode &t,
                              string indent = "") const {
    oss << indent << "|-- ";
    if (t->isTerminal()) {
        oss << t->token.word << endl;
    } else {
        oss << _symbols[t->ruleT.lhs] << endl;
        for (auto i : t->children) {
            printParseTree(oss, i, indent + "  ");
        }
    }
}

void MyParser::printParseTable(std::ostream &grammarOut) const {
    if(_parseTableHasConflict){
        cout << "[Warn] Parse table has conflict!!!!!\n";
        throw std::runtime_error("Parse table has conflict");
    }
    grammarOut << "state, ";
    for (int i = 1; i < _symbols.size(); ++i) {
        grammarOut << _symbols[i] << ", ";
    }
    grammarOut << endl;
    for (int i = 0; i < _parseTable.size(); ++i) {
        grammarOut << i << ",";
        for (int j = 1; j < _parseTable[i].size(); ++j) {
            grammarOut << _parseTable[i][j] << ", ";
        }
        grammarOut << endl;
    }
}

void MyParser::printSyntax(std::ostream &oss) const {
    for(auto i: _rules){
        oss << i.name << "    ,    ";
        for(auto j: i.rhs){
            oss << _symbols[j] << " ";
        }
        oss << "\n";
    }
}

MyParser::RuleType::RuleType(){};
MyParser::RuleType::RuleType(int index, string name) : index(index), name(name) {}
string MyParser::RuleType::getName() const { return name; }
bool MyParser::RuleType::operator==(const RuleType &t) const {
    return index == t.index && lhs == t.lhs && rhs == t.rhs;
}
bool MyParser::RuleType::operator!=(const RuleType &t) const { return !(*this == t); }
bool MyParser::RuleType::operator<(const RuleType &item) const {
    if (index != item.index) {
        return index < item.index;
    } else if (lhs != item.lhs) {
        return lhs < item.lhs;
    } else if (name != item.name) {
        return name < item.name;
    } else {
        return rhs < item.rhs;
    }
}


MyParser::LRItem::LRItem(MyParser::RuleType rule, int position, int lookahead)
    : rule(rule), position(position), lookahead(lookahead){};
int MyParser::LRItem::getExpectedSymbol() {
    if (position >= rule.rhs.size())
        return -1;
    return rule.rhs.at(position);
}
bool MyParser::LRItem::operator==(const LRItem &item) const {
    return rule == item.rule && position == item.position && lookahead == item.lookahead;
}
bool MyParser::LRItem::operator!=(const LRItem &item) const { return !(*this == item); }
bool MyParser::LRItem::operator<(const LRItem &item) const {
    if (position != item.position) {
        return position < item.position;
    } else if (lookahead != item.lookahead) {
        return lookahead < item.lookahead;
    } else {
        return rule < item.rule;
    }
}

MyParser::Action::Action() : atype(ActionType::UNDEFINED) {}
MyParser::Action::Action(bool) : atype(ActionType::ACCEPT) {}
MyParser::Action::Action(int state) : atype(ActionType::SHIFT), state(state) {}
MyParser::Action::Action(RuleType rule) : atype(ActionType::REDUCE), rule(rule) {}
void MyParser::Action::addAction(const Action &a) {
    if (atype != ActionType::CONFLICT)
        actions = vector<PtrAction>();
    actions.push_back(make_shared<MyParser::Action>(a));
}
std::ostream &operator<<(std::ostream &oss, const MyParser::Action &a) {
    switch (a.atype) {
    case MyParser::ActionType::UNDEFINED:
        oss << "-";
        break;
    case MyParser::ActionType::SHIFT:
        oss << "s" << a.state;
        break;
    case MyParser::ActionType::REDUCE:
        oss << "r" << a.rule.index;
        break;
    case MyParser::ActionType::ACCEPT:
        oss << "ACC";
        break;
    case MyParser::ActionType::CONFLICT:
        for (auto i : a.actions) {
            oss << i << "/";
        }
        oss << "\b";
        break;
    default:
        __PRINT_LINE_HINT("Miss a condition in switch");
        exit(0);
    }
    return oss;
}
bool MyParser::Action::operator==(const Action &a) const {
    if (atype != a.atype)
        return false;
    switch (atype) {
    case ActionType::ACCEPT:
    case ActionType::UNDEFINED:
        return true;
    case ActionType::SHIFT:
        return state == a.state;
    case ActionType::REDUCE:
        return rule == a.rule;
    case ActionType::CONFLICT:
    case ActionType::ERROR:
        return false;
    default:
        __PRINT_LINE_HINT("Miss a condition in switch");
        exit(0);
    }
}
bool MyParser::Action::operator!=(const Action &b) const { return !(*this == b); }

void MyParser::Action::copyFields(const Action &a) {
    atype = a.atype;
    switch (atype) {
    case ActionType::UNDEFINED:
    case ActionType::ACCEPT:
        break;
    case ActionType::REDUCE:
        rule = RuleType();
        rule = a.rule;
        break;
    case ActionType::SHIFT:
        state = a.state;
        break;
    case ActionType::CONFLICT:
        actions = vector<PtrAction>();
        actions = a.actions;
        break;
    default: // should never reach
        __PRINT_LINE_HINT("Miss a condition in switch");
        exit(0);
    }
}

bool MyParser::ParseTreeNode::isTerminal() { return children.empty(); }
MyParser::ParseTreeNode::ParseTreeNode() {}
MyParser::ParseTreeNode::ParseTreeNode(RuleType ruleT, deque<PtrParseTreeNode> children)
    : ruleT(ruleT), children(children){};
MyParser::ParseTreeNode::ParseTreeNode(Token token) : token(token){};
MyParser::ParseTreeNode::~ParseTreeNode() {}
